package com.lwl.shua.class03;

import com.lwl.Utils.PrintUtil;
import org.junit.Test;

import java.util.*;

/**
 * @author lwl
 * @Description 给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 k ，返回到目标结点 target 距离为 k 的所有结点的值的数组。
 * <p>
 * 答案可以以 任何顺序 返回。
 * @date 2025/9/7 16:52
 */
public class Code09Lc863DistanceK {
    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }
    }

    @Test
    public void test() {
        // [0,2,1,null,null,3]
        TreeNode root = new TreeNode(0);
        root.left = new TreeNode(2);
        root.right = new TreeNode(1);
        root.right.left = new TreeNode(3);
        TreeNode target = root.right.left;
        List<Integer> res = distanceK(root, target, 3);
        PrintUtil.printList(res);
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        // 方法 1
        // 从 target 节点出发往下寻找，找到距离为 k 节点集合
        // 从 root 节点出发往下寻找，找到距离为 k - （root - target） 节点集合
        // 方法 2
        // 记录下每个节点的父节点，则题目将变形为图的宽度优先遍历，
        Map<TreeNode, TreeNode> parentMap = createParentMap(root);
        LinkedList<TreeNode> queue = new LinkedList<>();
        Set<TreeNode> set = new HashSet<>();
        queue.addLast(target);
        set.add(target);
        int curDistance = 0;
        TreeNode curLevelLastNode = target;
        TreeNode nextLevelLastNode = null;
        List<Integer> res = new ArrayList<>();
        while (!queue.isEmpty()) {
            TreeNode node = queue.pollFirst();
            if (curDistance == k) {
                res.add(node.val);
            }
            if (node.left != null && !set.contains(node.left)) {
                queue.addLast(node.left);
                set.add(node.left);
                nextLevelLastNode = node.left;
            }
            if (node.right != null && !set.contains(node.right)) {
                queue.addLast(node.right);
                set.add(node.right);
                nextLevelLastNode = node.right;
            }
            if (parentMap.get(node) != null && !set.contains(parentMap.get(node))) {
                queue.addLast(parentMap.get(node));
                set.add(parentMap.get(node));
                nextLevelLastNode = parentMap.get(node);
            }
            if (curLevelLastNode == node) {
                curDistance++;
                curLevelLastNode = nextLevelLastNode;
            }
        }
        return res;
    }

    private Map<TreeNode, TreeNode> createParentMap(TreeNode root) {
        Map<TreeNode, TreeNode> parentMap = new HashMap<>();
        LinkedList<TreeNode> queue = new LinkedList<>();
        queue.addLast(root);
        parentMap.put(root, null);
        while (!queue.isEmpty()) {
            TreeNode node = queue.pollFirst();
            if (node.left != null) {
                queue.addLast(node.left);
                parentMap.put(node.left, node);
            }
            if (node.right != null) {
                queue.addLast(node.right);
                parentMap.put(node.right, node);
            }
        }
        return parentMap;
    }
}
